히든 스테이트를 머신러닝 모델로 구성 ["Our key idea is to make the hidden state itself a model f with weights W", Section 2.1, p.4]

• TTT Layer는 히든 스테이트를 머신러닝 모델(f)로 구성하고, 파라미터(W)를 통해 컨텍스트 정보를 저장함 ["The hidden state st is now equivalent to Wt, the weights of a model f", Section 2.1, p.4]

• 모든 시퀀스 데이터에 대해 테스트 타임에 학습이 이뤄짐으로써 각 시퀀스에 맞는 최적의 파라미터를 찾음 ["Even at test time, our new layer still trains a different sequence of weights W1,...,WT for every input sequence", Section 2.1, p.5]

• Inner loop(TTT)와 Outer loop(네트워크 학습)의 이중 학습 구조를 가짐 ["We refer to training the larger network as the outer loop, and training W within each TTT layer as the inner loop", Section 2.2, p.6]

업데이트 규칙을 자기지도학습 단계로 설정 ["and the update rule a step of self-supervised learning", Section 2.1, p.4]

- 입력 토큰 xt를 corrupted input x̃t로 변환하고 이를 다시 복원하는 task를 학습 ["One choice of ℓ is reconstructing xt itself. To make the learning problem nontrivial, we first process xt into a corrupted input x̃t", Section 2.3, p.6-7]

- θK, θV, θQ 세 개의 학습 가능한 projection matrices를 도입:
	- Input Token xt가 들어오면 3가지 뷰로 투영되는데, 이 중 Model Weights (Wt)는 Training view와 Label view를 사용해 self-supervised learning으로 업데이트 / 최종 Output (zt) 생성은 업데이트된 weight로 Test View를 처리하여 진행
	- <svg viewBox="0 0 800 500" xmlns="http://www.w3.org/2000/svg">

출력함수 f에 대해 TTT-Linear와 TTT-MLP 두 가지 구현 제안 ["We propose two variants of TTT layers – TTT-Linear and TTT-MLP", Section 2.7, p.12]

• TTT-Linear: 히든 스테이트가 선형 모델(f(x) = Wx)로 구성됨

• TTT-MLP: 히든 스테이트가 2층 MLP로 구성됨, 더 표현력이 높음 (2-layer MLP with GELU activation)

• 둘 다 Layer Normalization과 residual connection을 포함 ["For TTT-Linear, flin(x) = Wx, where W is square. For TTT-MLP, fMLP has two layers similar to the MLPs in Transformers", Section 2.7, p.12]

Mini-batch TTT를 통한 병렬화 ["Our proposed solution – mini-batch gradient descent", Section 2.4, p.7]

• 기존 online gradient descent는 순차적 의존성으로 인해 병렬화가 어려움
• Mini-batch TTT는 batch size b만큼의 토큰을 모아서 한 번에 그래디언트 계산
• 각 미니배치 내의 그래디언트 계산이 독립적이므로 병렬 처리 가능 ["To parallelize Gt for t = 1,...,T, we can take all of them w.r.t. W0... Our proposed solution – mini-batch gradient descent", Section 2.4, p.7-8]

Dual form을 통한 하드웨어 효율성 향상 ["We call this procedure the dual form, in contrast to the primal form", Section 2.5, p.7-8]

• 하드웨어 효율성을 위해 중간 변수들(Gs, Ws)를 실제로 계산하지 않음
• Matrix multiplication을 주로 사용하여 최종 출력(Wb, z1,...,zb)만 직접 계산
• 현대 하드웨어(GPU, TPU)의 matrix multiplication 최적화 활용 ["The goal of the dual form is to compute z̄XK+1 and W1b,...,WKb with only matmuls and light-weight operations such as sums, σ, and σ'", Section 2.5, p.8]

• ***dual form은 primal form에 비해:
  • 이론적 복잡도는 약간 높지만(O(b×d²) vs O(b²×d))
  • 하드웨어 활용도가 높아 실제로는 5배 이상 빠름 ["In our JAX implementation, training with the dual form is more than 5× faster than with primal", Section 2.5, p.9]

메모리 효율적인 그래디언트 체크포인팅 ["A standard technique to save memory in this scenario is gradient checkpointing", Appendix C, p.31]

• 일반적으로 JAX나 PyTorch같은 라이브러리들은 backward pass를 위해 forward pass의 중간 결과(activation)들을 모두 저장함 ["By default, libraries such as JAX and PyTorch save the intermediate activations during a forward pass so they can be reused during the backward pass", Appendix C, p.31]
• TTT layer의 경우 히든 스테이트 W1,...,WT를 모두 저장해야 하므로 메모리 사용량이 매우 큼

해결 방법:

1. 시간 축으로의 체크포인팅:

   • Mini-batch TTT에서는 κ = T/b 개의 체크포인트만 저장 (b는 batch size)
   • 각 미니배치의 마지막 W만 저장하고 나머지는 필요할 때 다시 계산 ["With TTT mini-batch and the dual form, we still need to save (assume divisible) κ = T/b Ws at the end of the mini-batches", Appendix C, p.31]
2. 메모리-계산 트레이드오프:

   • 메모리 사용량은 줄어들지만 일부 계산을 다시 해야 함
   • 기존의 layer 방향 체크포인팅과 달리 시간 방향으로 적용 ["A standard technique to save memory in this scenario is gradient checkpointing, which is usually applied through layers, but we apply it through time", Appendix C, p.31]

장점:

1. 메모리 효율성:

   • W의 전체 시퀀스 대신 체크포인트들만 저장
   • 긴 시퀀스 처리가 가능해짐
2. 유연성:

   • Mini-batch size를 조절하여 메모리-속도 트레이드오프 조정 가능
   • 하드웨어 제약에 맞춰 최적화 가능

   *** 특히 긴 시퀀스를 처리할 때, 일반적인 layer 방향의 체크포인팅과 달리 시간 축으로 적용되어 TTT의 특성을 잘 활용할 수 있게 됨